import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
import sys, select, termios, tty
import time


# 获取键盘输入的函数     
def getKey():
    # 将终端设置为原始模式，这样可以直接读取单个字符，而不需要按回车键
    tty.setraw(sys.stdin.fileno())
    # 使用 select 函数监听标准输入，设置超时时间为 0.1 秒
    rlist, _, _ = select.select([sys.stdin], [], [], 0.1)
    if rlist:
        # 如果有输入，读取一个字符
        key = sys.stdin.read(1)
    else:
        # 如果没有输入，将 key 设为空字符串
        key = ''

    # 恢复终端的原始属性，避免影响后续输入
    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)
    return key


# 键盘控制节点类，继承自 rclpy 的 Node 类
class KeyboardControlNode(Node):
    def __init__(self):
        # 调用父类的构造函数，初始化节点，节点名为 'keyboard_control_node'
        super().__init__('keyboard_control_node')
        # 创建一个发布者，用于发布 Twist 类型的消息到 'twist_cmd' 话题，队列大小为 10
        self.publisher_ = self.create_publisher(Twist, 'twist_cmd', 10)
        # 初始化线速度为 0
        self.linear_vel = 0.0
        # 初始化角速度为 0
        self.angular_vel = 0.0
        # 设置线速度每次变化的步长为 0.1
        self.linear_step = 0.005
        # 设置角速度每次变化的步长为 500
        self.angular_step = 0.2
        # 打印信息提示用户可以使用 w/s/a/d 键来控制
        self.get_logger().info('Keyboard control node started. Use w/s/q/e/a/d keys to control.')
        # 用于记录按键的状态，键为按键字符，值为布尔类型，表示按键是否被按下
        self.key_states = {}

    # 根据按键更新速度的函数
    def update_velocity(self, key):
        if key == 'w':
            # 记录 'w' 键被按下
            self.key_states[key] = True
            # 线速度增加一个步长
            self.linear_vel += self.linear_step
        elif key =='s':
            # 记录's' 键被按下
            self.key_states[key] = True
            # 线速度减少一个步长
            self.linear_vel -= self.linear_step
        elif key == 'q':
            self.key_states[key] = True
            self.linear_vel = 300.0
            self.angular_vel += self.angular_step
        elif key == 'e':
            self.key_states[key] = True
            self.linear_vel = 300.0
            self.angular_vel -= self.angular_step
        elif key == 'a':
            self.key_states[key] = True
            self.linear_vel = -300.0
            self.angular_vel += self.angular_step
        elif key == 'd':
            self.key_states[key] = True
            self.linear_vel = -300.0
            self.angular_vel -= self.angular_step

        # 限制线速度在 -500.0 到 500.0 的范围内
        self.linear_vel = max(-0.4, min(0.4, self.linear_vel))
        # 限制角速度在 -2500.0 到 2500.0 的范围内
        self.angular_vel = max(-2.0, min(2.0, self.angular_vel))

    # 当按键松开时，将对应的速度设置为 0 的函数
    def release_key(self, key):
        if key in self.key_states:
            # 标记该按键为松开状态
            self.key_states[key] = False
            if key in ['w','s', 'q', 'e', 'a', 'd']:
                if key in ['w','s']:
                    self.linear_vel = 0.0
                elif key in ['q', 'a']:
                    self.linear_vel = 0.0
                    self.angular_vel = 0.0
                elif key in ['e', 'd']:
                    self.linear_vel = 0.0
                    self.angular_vel = 0.0

    # 发布速度消息的函数
    def publish_velocity(self):
        # 创建一个 Twist 类型的消息对象
        msg = Twist()
        # 设置线速度的 x 分量为当前的线速度
        msg.linear.x = self.linear_vel
        # 设置角速度的 z 分量为当前的角速度
        msg.angular.z = self.angular_vel
        # 打印当前的线速度和角速度
        self.get_logger().info(f"当前线速度: {self.linear_vel}, 当前角速度: {self.angular_vel}")
        # 发布消息
        self.publisher_.publish(msg)
        time.sleep(0.1)


def main(args=None):
    # 初始化 ROS 2
    rclpy.init(args=args)
    # 创建键盘控制节点的实例
    node = KeyboardControlNode()

    global settings
    # 保存当前终端的属性
    settings = termios.tcgetattr(sys.stdin)
    try:
        while rclpy.ok():
            # 获取键盘输入
            key = getKey()
            if key:
                # 如果有按键输入，更新速度
                node.update_velocity(key)
            else:
                # 如果没有按键输入，检查是否有按键松开
                for k in list(node.key_states.keys()):
                    if node.key_states[k]:
                        # 如果按键处于按下状态，将其视为松开并处理
                        node.release_key(k)
            # 发布当前的速度消息
            node.publish_velocity()
            # 按照设定的频率暂停线程

    except Exception as e:
        # 捕获并打印异常信息
        print(e)
    finally:
        # 恢复终端的原始属性
        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)

    # 销毁节点
    node.destroy_node()
    # 关闭 ROS 2
    rclpy.shutdown()


if __name__ == '__main__':
    main()